#include "VerifyCore.h"

// 全局变量
STATIC CHAR16 gAuthListPath[] = L"\\EFI\\BOOT\\Auth.list";
STATIC CHAR16 gCoreEfiPath[] = L"\\EFI\\BOOT\\Core.efi";

// 从文件中读取指定长度的内容
STATIC
EFI_STATUS
ReadFileContent (
  IN  CHAR16  *FilePath,
  IN  UINTN   Offset,
  OUT VOID    *Buffer,
  IN  UINTN   Length
  )
{
  EFI_STATUS                      Status;
  EFI_SIMPLE_FILE_SYSTEM_PROTOCOL *SimpleFs;
  EFI_FILE_PROTOCOL              *Root;
  EFI_FILE_PROTOCOL              *File;
  UINTN                          ReadSize;

  // Try all handles that implement Simple File System. Some environments
  // register multiple file systems and LocateProtocol may return a wrong
  // instance. Iterate handles and try to open the requested file on each
  // volume until we succeed.
  EFI_HANDLE *HandleBuffer = NULL;
  UINTN HandleCount = 0;
  Status = gBS->LocateHandleBuffer(
                  ByProtocol,
                  &gEfiSimpleFileSystemProtocolGuid,
                  NULL,
                  &HandleCount,
                  &HandleBuffer
                  );
  if (EFI_ERROR(Status) || HandleCount == 0) {
    Print(L"ReadFileContent: LocateHandleBuffer failed: %r\n", Status);
    return Status;
  }

  // Prepare an alternate path with a leading '\\' if callers passed a
  // path without it (e.g. L"EFI\\BOOT\\Core.efi" vs L"\\EFI\\BOOT\\Core.efi").
  BOOLEAN NeedFreeAlt = FALSE;
  CHAR16 *AltPath = NULL;
  if (FilePath[0] != L'\\') {
    UINTN Len = StrLen(FilePath) + 2; // leading '\\' + NULL
    Status = gBS->AllocatePool(EfiBootServicesData, Len * sizeof(CHAR16), (VOID**)&AltPath);
    if (!EFI_ERROR(Status)) {
      NeedFreeAlt = TRUE;
      AltPath[0] = L'\\';
      StrCpyS(AltPath + 1, Len - 1, FilePath);
    }
  }

  UINTN i;
  EFI_STATUS OpenStatus = EFI_NOT_FOUND;
  for (i = 0; i < HandleCount; i++) {
    SimpleFs = NULL;
    Status = gBS->HandleProtocol(HandleBuffer[i], &gEfiSimpleFileSystemProtocolGuid, (VOID**)&SimpleFs);
    if (EFI_ERROR(Status) || SimpleFs == NULL) {
      continue;
    }

    Status = SimpleFs->OpenVolume(SimpleFs, &Root);
    if (EFI_ERROR(Status) || Root == NULL) {
      continue;
    }

    // Try provided path first
    Status = Root->Open(Root, &File, FilePath, EFI_FILE_MODE_READ, 0);
    if (EFI_ERROR(Status) && AltPath != NULL) {
      // Try alternate path with leading '\\'
      Status = Root->Open(Root, &File, AltPath, EFI_FILE_MODE_READ, 0);
    }

    if (!EFI_ERROR(Status)) {
      OpenStatus = Status;
      break;
    }

    // Not found on this volume, close and continue
    Root->Close(Root);
  }

  if (HandleBuffer != NULL) {
    FreePool(HandleBuffer);
  }
  if (NeedFreeAlt && AltPath != NULL) {
    FreePool(AltPath);
  }

  if (EFI_ERROR(OpenStatus)) {
    Print(L"ReadFileContent: Root->Open failed for %s: %r\n", FilePath, OpenStatus);
    return OpenStatus;
  }

  Status = File->SetPosition(File, Offset);
  if (EFI_ERROR(Status)) {
    File->Close(File);
    Root->Close(Root);
    Print(L"ReadFileContent: SetPosition failed: %r\n", Status);
    return Status;
  }

  ReadSize = Length;
  Status = File->Read(File, &ReadSize, Buffer);
  
  // 打印实际读取情况
  if (EFI_ERROR(Status)) {
    Print(L"ReadFileContent: File->Read failed: %r (requested=%u, read=%u)\n", Status, Length, ReadSize);
  }

  File->Close(File);
  Root->Close(Root);

  return Status;
}

// 写入指定长度的内容到文件
STATIC
EFI_STATUS
WriteFileContent (
  IN  CHAR16  *FilePath,
  IN  UINTN   Offset,
  IN  VOID    *Buffer,
  IN  UINTN   Length
  )
{
  EFI_STATUS                      Status;
  EFI_SIMPLE_FILE_SYSTEM_PROTOCOL *SimpleFs;
  EFI_FILE_PROTOCOL              *Root;
  EFI_FILE_PROTOCOL              *File;

  // Iterate all Simple File System handles and try to open the file for
  // read/write. Also try an alternate path with a leading '\\' if needed.
  EFI_HANDLE *HandleBuffer = NULL;
  UINTN HandleCount = 0;
  Status = gBS->LocateHandleBuffer(
                  ByProtocol,
                  &gEfiSimpleFileSystemProtocolGuid,
                  NULL,
                  &HandleCount,
                  &HandleBuffer
                  );
  if (EFI_ERROR(Status) || HandleCount == 0) {
    return Status;
  }

  BOOLEAN NeedFreeAlt = FALSE;
  CHAR16 *AltPath = NULL;
  if (FilePath[0] != L'\\') {
    UINTN Len = StrLen(FilePath) + 2;
    Status = gBS->AllocatePool(EfiBootServicesData, Len * sizeof(CHAR16), (VOID**)&AltPath);
    if (!EFI_ERROR(Status)) {
      NeedFreeAlt = TRUE;
      AltPath[0] = L'\\';
      StrCpyS(AltPath + 1, Len - 1, FilePath);
    }
  }

  EFI_STATUS OpenStatus = EFI_NOT_FOUND;
  for (UINTN i = 0; i < HandleCount; i++) {
    SimpleFs = NULL;
    Status = gBS->HandleProtocol(HandleBuffer[i], &gEfiSimpleFileSystemProtocolGuid, (VOID**)&SimpleFs);
    if (EFI_ERROR(Status) || SimpleFs == NULL) {
      continue;
    }

    Status = SimpleFs->OpenVolume(SimpleFs, &Root);
    if (EFI_ERROR(Status) || Root == NULL) {
      continue;
    }

    Status = Root->Open(Root, &File, FilePath, EFI_FILE_MODE_READ | EFI_FILE_MODE_WRITE, 0);
    if (EFI_ERROR(Status) && AltPath != NULL) {
      Status = Root->Open(Root, &File, AltPath, EFI_FILE_MODE_READ | EFI_FILE_MODE_WRITE, 0);
    }

    if (!EFI_ERROR(Status)) {
      OpenStatus = Status;
      break;
    }

    Root->Close(Root);
  }

  if (HandleBuffer != NULL) {
    FreePool(HandleBuffer);
  }
  if (NeedFreeAlt && AltPath != NULL) {
    FreePool(AltPath);
  }

  if (EFI_ERROR(OpenStatus)) {
    return OpenStatus;
  }

  Status = File->SetPosition(File, Offset);
  if (EFI_ERROR(Status)) {
    File->Close(File);
    Root->Close(Root);
    return Status;
  }

  Status = File->Write(File, &Length, Buffer);
  
  File->Close(File);
  Root->Close(Root);
  
  return Status;
}

// 解析形如 "<30d>", "<12h>", "<1d>" 的到期字符串，返回秒数
STATIC
UINT64
ParseExpirationString(
  IN CHAR8 *Str,
  IN UINTN Len
  )
{
  if (Str == NULL || Len == 0) {
    return 0;
  }

  UINTN i = 0;
  // 跳过空格
  while (i < Len && Str[i] == ' ') {
    i++;
  }
  if (i >= Len || Str[i] != '<') {
    return 0;
  }
  i++;

  UINT64 num = 0;
  while (i < Len && Str[i] >= '0' && Str[i] <= '9') {
    num = num * 10 + (Str[i] - '0');
    i++;
  }

  if (i >= Len) {
    return 0;
  }

  CHAR8 unit = Str[i];
  UINT64 seconds = 0;
  if (unit == 'd') {
    seconds = num * 24ULL * 60ULL * 60ULL;
  } else if (unit == 'h') {
    seconds = num * 60ULL * 60ULL;
  } else if (unit == 'm') {
    seconds = num * 60ULL;
  } else {
		seconds = 0;
	}

  return seconds;
}

// 将 EFI_TIME 转为秒数的简易函数（近似，按天数计算）
STATIC
UINT64
EfiTimeToSeconds(
  IN EFI_TIME *T
  )
{
  if (T == NULL) {
    return 0;
  }
  UINT64 secs = 0;
  secs += (UINT64)T->Year * 365ULL * 24ULL * 60ULL * 60ULL;
  secs += (UINT64)T->Month * 30ULL * 24ULL * 60ULL * 60ULL;
  secs += (UINT64)T->Day * 24ULL * 60ULL * 60ULL;
  secs += (UINT64)T->Hour * 60ULL * 60ULL;
  secs += (UINT64)T->Minute * 60ULL;
  secs += (UINT64)T->Second;
  return secs;
}

// 验证GUID和MAC地址
EFI_STATUS
VerifyGuidAndMac (
  VOID
  )
{
  EFI_STATUS  Status;
  CHAR8       AuthGuid[GUID_LENGTH + 1];
  CHAR8       CoreGuid[GUID_LENGTH + 1];
  CHAR8       AuthMac[MAC_LENGTH + 1];
  CHAR8       CoreMac[MAC_LENGTH + 1];
  UINTN       LineCount;
  CHAR8       *Line;
  UINTN       LineLength;
  
  // 读取Core.efi中的GUID
  Status = ReadFileContent(
             gCoreEfiPath,
             CORE_GUID_OFFSET,
             CoreGuid,
             GUID_LENGTH
             );
  if (EFI_ERROR(Status)) {
    Print(L"Failed to read GUID from Core.efi\n");
    return VERIFY_ERROR_READ_FAILED;
  }
  CoreGuid[GUID_LENGTH] = '\0';

  // 读取Auth.list第一行的GUID
  Status = ReadFileContent(
             gAuthListPath,
             0,
             AuthGuid,
             GUID_LENGTH
             );
  if (EFI_ERROR(Status)) {
    Print(L"Failed to read GUID from Auth.list\n");
    return VERIFY_ERROR_READ_FAILED;
  }
  AuthGuid[GUID_LENGTH] = '\0';

  // 比较GUID
  if (AsciiStrCmp(AuthGuid, CoreGuid) != 0) {
    Print(L"GUID mismatch\n");
    return VERIFY_ERROR_GUID_MISMATCH;
  }

  // // 读取Core.efi中的MAC
  // Status = ReadFileContent(
  //            gCoreEfiPath,
  //            CORE_MAC_OFFSET,
  //            CoreMac,
  //            MAC_LENGTH
  //            );
  // if (EFI_ERROR(Status)) {
  //   Print(L"Failed to read MAC from Core.efi\n");
  //   return VERIFY_ERROR_READ_FAILED;
  // }
  // CoreMac[MAC_LENGTH] = '\0';


  // // 读取Auth.list第二行的MAC
  // Status = ReadFileContent(
  //            gAuthListPath,
  //            GUID_LENGTH + 2, // +2 for newline
  //            AuthMac,
  //            MAC_LENGTH
  //            );
  // if (EFI_ERROR(Status)) {
  //   Print(L"Failed to read MAC from Auth.list\n");
  //   return VERIFY_ERROR_READ_FAILED;
  // }
  // AuthMac[MAC_LENGTH] = '\0';


  // // 比较MAC
  // if (AsciiStrCmp(AuthMac, CoreMac) != 0) {
  //   return VERIFY_ERROR_MAC_MISMATCH;
  // }

  return VERIFY_SUCCESS;
}

// EFI_TIME
// 处理首次加载和到期时间
EFI_STATUS
HandleFirstTimeAndExpiration (
  VOID
  )
{
	UINTN StartOffset = CORE_FIRST_TIME_OFFSET;
	VOID* FirstTimeBuffer = NULL;
	gBS->AllocatePool(EfiBootServicesData, sizeof(EFI_TIME) + 1, &FirstTimeBuffer);
	if (FirstTimeBuffer == NULL) {
		return EFI_OUT_OF_RESOURCES;
	}

	EFI_STATUS Status = ReadFileContent(
		gCoreEfiPath,
		StartOffset,
		FirstTimeBuffer,	
		sizeof(EFI_TIME)
	);
	if (EFI_ERROR(Status)) {
		FreePool(FirstTimeBuffer);
		Print(L"Failed to read first time from Core.efi\n");	
		return VERIFY_ERROR_READ_FAILED;
	}

	// 检查是否为第一次加载
	EFI_TIME* FirstTime = (EFI_TIME*)FirstTimeBuffer;
	BOOLEAN IsFirstLoad = (FirstTime->Year == 0 && FirstTime->Month == 0 && FirstTime->Day == 0);
	if (IsFirstLoad) {
		// 获取当前时间
		EFI_TIME CurrentTime = { 0 };
		Status = gST->RuntimeServices->GetTime(&CurrentTime, NULL);
		if (EFI_ERROR(Status)) {
			FreePool(FirstTimeBuffer);
			Print(L"Failed to get current time\n");
			return VERIFY_ERROR_READ_FAILED;
		}

		// 写入当前时间作为首次加载时间
		Status = WriteFileContent(
			gCoreEfiPath,
			StartOffset,
			&CurrentTime,
			sizeof(EFI_TIME)
		);

    // 写入nvram

		if (EFI_ERROR(Status)) {
			FreePool(FirstTimeBuffer);
			Print(L"Failed to write first time to Core.efi\n");
			return VERIFY_ERROR_READ_FAILED;
		}

		FreePool(FirstTimeBuffer);
		return VERIFY_SUCCESS;
	}

	// 如果不是第一次加载
	StartOffset = CORE_TIME_OFFSET;
	VOID* ExpireTimeBuffer;
	gBS->AllocatePool(EfiBootServicesData, 6, &ExpireTimeBuffer);
	if (ExpireTimeBuffer == NULL) {
		FreePool(FirstTimeBuffer);	
		return EFI_OUT_OF_RESOURCES;
	}

	Status = ReadFileContent(
		gCoreEfiPath,	
		StartOffset,
		ExpireTimeBuffer,
		6
	);
	if (EFI_ERROR(Status)) {	
		FreePool(ExpireTimeBuffer);
		Print(L"Failed to read expire time from Core.efi\n");
		return VERIFY_ERROR_READ_FAILED;
	}

  // 解析到期时间并检查是否过期
	UINT64 ExpireSecondsParsed = ParseExpirationString((CHAR8*)ExpireTimeBuffer, 6);
	if (ExpireSecondsParsed == 0) {
			Print(L"Failed to parse expiration string: %a\n", (CHAR8*)ExpireTimeBuffer);
			FreePool(ExpireTimeBuffer);
			FreePool(FirstTimeBuffer);
			return VERIFY_ERROR_READ_FAILED;
	}

	// 将 FirstTime 转换为秒
	UINT64 FirstLoadSecs = EfiTimeToSeconds(FirstTime);

	// 获取当前时间并转换为秒
	EFI_TIME CurTime = {0};
	Status = gST->RuntimeServices->GetTime(&CurTime, NULL);
	if (EFI_ERROR(Status)) {
			FreePool(ExpireTimeBuffer);
			FreePool(FirstTimeBuffer);
			Print(L"Failed to get current time for expiration check\n");
			return VERIFY_ERROR_READ_FAILED;
	}
	UINT64 CurrentSecs = EfiTimeToSeconds(&CurTime);

	if (FirstLoadSecs + ExpireSecondsParsed < CurrentSecs) {
			FreePool(ExpireTimeBuffer);
			FreePool(FirstTimeBuffer);
			return VERIFY_ERROR_TIME_EXPIRED;
	}

	FreePool(ExpireTimeBuffer);
	FreePool(FirstTimeBuffer);

	return VERIFY_SUCCESS;
}

// 主入口函数
// EFI_STATUS
// EFIAPI
// UefiMain (
//   IN EFI_HANDLE        ImageHandle,
//   IN EFI_SYSTEM_TABLE  *SystemTable
//   )
// {
//   EFI_STATUS  Status;

//   // 验证GUID和MAC
//   Status = VerifyGuidAndMac();
//   if (EFI_ERROR(Status)) {
//     return Status;
//   }

// 	if (Status == VERIFY_ERROR_GUID_MISMATCH || Status == VERIFY_ERROR_MAC_MISMATCH) {
// 		Print(L"GUID or MAC address mismatch!\n");
// 		return Status;
// 	}

//   // 处理首次加载和到期时间
//   Status = HandleFirstTimeAndExpiration();
//   if (EFI_ERROR(Status)) {
//     return Status;
//   }

// 	if (Status == VERIFY_ERROR_TIME_EXPIRED) {
// 		Print(L"License has expired!\n");
// 		return Status;
// 	}

//   Print(L"Verification completed successfully.\n");
//   return EFI_SUCCESS;
// }